"""
各种优化器类
"""

import numpy as np

class Optimizer:
    """优化器基类"""
    def __init__(self):
        self.target = None
        self.hooks = []

    def setup(self,target):
        """将Layer的参数变量设置为target实例变量"""
        self.target = target
        return self

    def update(self):
        """对除grad=None之外的参数进行更新"""
        params = [p for p in self.target.params() if p.grad is not None]
        #预处理(可选)
        for f in self.hooks:
            f(params)
        #更新参数
        for param in params:
            self.update_one(param)

    def update_one(self, param):
        """具体的参数更新方法"""
        raise NotImplementedError()
    def add_hook(self, f):
        """
        用户可使用add_hook方法添加在更新前进行预处理的函数
        此机制可用于权重衰减、梯度裁剪等
        """
        self.hooks.append(f)


class SGD(Optimizer):
    """随机梯度下降法"""
    def __init__(self, lr=0.01):
        """
        SGD随机梯度下降法
        参数：lr--学习率
        """
        super().__init__()
        self.lr = lr
    def update_one(self, param):
        """
        梯度下降法更新
        W := W - lr * grad
        """
        param.data -= self.lr * param.grad.data

class MomentumSGD(Optimizer):
    """MomentumSGD 动量下降法"""
    def __init__(self, lr=0.01, momentum=0.9):
        """
        SGD随机梯度下降法
        参数：lr--学习率, momentum--动量系数
        """
        super().__init__()
        self.lr = lr
        self.momentum = momentum
        self.vs = {} #空字典
    def update_one(self, param):
        """
        动量梯度下降法更新
        v := momentum * v - lr * grad
        W := W + v
        """
        v_key = id(param)
        if v_key not in self.vs:
            self.vs[v_key] = np.zeros_like(param.data)
        v = self.vs[v_key]
        v *= self.momentum
        v -= self.lr * param.grad.data
        param.data += v








